package com.jxpanda.spring.module.auth.core.token;

import com.jxpanda.spring.module.auth.config.properties.JwtProperties;
import com.jxpanda.spring.module.auth.core.authentication.token.UserDetailsAuthenticationToken;
import lombok.Builder;
import org.springframework.lang.Nullable;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.security.oauth2.jwt.JwsHeader;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.security.oauth2.jwt.JwtClaimsSet;
import org.springframework.security.oauth2.jwt.JwtEncoder;
import reactor.core.publisher.Mono;

import java.time.Duration;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.Objects;
import java.util.UUID;

@Builder
public class JwtReactiveOAuth2TokenGenerator implements ReactiveOAuth2TokenGenerator<Jwt> {

    private final int expiresIn;

    private final JwtEncoder jwtEncoder;

    private final JwtProperties jwtProperties;

    private final ReactiveJwtCustomizer reactiveJwtCustomizer;

    private JwtReactiveOAuth2TokenGenerator(Integer expiresIn, JwtEncoder jwtEncoder, JwtProperties jwtProperties, @Nullable ReactiveJwtCustomizer reactiveJwtCustomizer) {
        this.expiresIn = Objects.requireNonNull(expiresIn);
        this.jwtEncoder = Objects.requireNonNull(jwtEncoder);
        this.jwtProperties = Objects.requireNonNull(jwtProperties);
        this.reactiveJwtCustomizer = Objects.requireNonNullElse(reactiveJwtCustomizer, new ReactiveJwtCustomizer() {
        });
    }

    @Override
    public <U extends UserDetails> Mono<Jwt> generate(UserDetailsAuthenticationToken<U> userDetailsAuthenticationToken) {
        return buildJwtEncodingContext(userDetailsAuthenticationToken)
                .map(this.reactiveJwtCustomizer::buildEncoderParameters)
                .map(this.jwtEncoder::encode);
    }

    private <U extends UserDetails> Mono<ReactiveJwtEncodingContext<U>> buildJwtEncodingContext(UserDetailsAuthenticationToken<U> userDetailsAuthenticationToken) {
        U userDetails = userDetailsAuthenticationToken.getUserDetails();
        JwsHeader.Builder headerBuilder = JwsHeader.with(jwtProperties.getAlgorithm().getJwsAlgorithm());
        Instant issuedAt = Instant.now();
        Instant expiresAt = issuedAt.plus(Duration.of(expiresIn, ChronoUnit.SECONDS));
        JwtClaimsSet.Builder claimsBuilder = JwtClaimsSet.builder()
                .issuer(jwtProperties.getIssuer())
                .subject(userDetails.getUsername())
                .audience(jwtProperties.getAudience())
                .issuedAt(issuedAt)
                .notBefore(issuedAt)
                .expiresAt(expiresAt)
                .claim("scope", userDetails.getAuthorities().stream().map(GrantedAuthority::getAuthority).toList())
                .id(UUID.randomUUID().toString());
        return Mono.just(new ReactiveJwtEncodingContext<>(headerBuilder, claimsBuilder, userDetailsAuthenticationToken));
    }

}
